#include "DeBruijnGraphBuilder.h" // include Global.h first
#include <iostream>
#include <iomanip>
#include <ctime>
#include <cstdlib>
#include <fstream>
#include <cmath>
#include "SequenceGraph.h"
#include "GraphErrorCorrector.h"
#include "FibHeap.h"
#include "Disambiguation.h"
#include "KmerCounter.h"
#include "GraphExaminator.h"
#include "SequenceGraphViz.h"

#if _WIN32
#include "XGetOpt.h"
#else
#include <unistd.h>
#endif

struct ConfigArgs {
	friend inline ostream & operator << (ostream& os, const ConfigArgs& confArgs) {
		os <<"kmerLength: " <<confArgs.kmerLength <<endl;
		os <<"kmerThreshold: " <<confArgs.kmerThreshold <<endl;
		os <<"lengthThreshold: " <<confArgs.lengthThreshold <<endl;
		os <<"multiplicityThreshold: " << confArgs.multiplicityThreshold <<endl;
		os <<"lowCoverage: " <<confArgs.lowCoverage <<endl;
		os <<"insertLength: " << confArgs.insertLength <<endl;
		os <<"insertLenSD: " << confArgs.insertLenSD << endl;
		os <<"distBound: " << confArgs.distBound << endl;
		os <<"pairSupport: " << confArgs.pairSupport << endl;
		os <<"resampling: " << confArgs.resampling << endl;
		os <<"inward: " << confArgs.inward <<endl;
		os <<"largeVertexSize: " << confArgs.largeVertexSize <<endl;
		os <<"Read file location: " << confArgs.inputName <<endl;
		os <<"Output file location: " << confArgs.outputName <<endl;
		return os;
	}

	int kmerLength;
	int kmerThreshold;
	double multiplicityThreshold;
	int lengthThreshold;
	double lowCoverage;
	int insertLength;
	int insertLenSD;
	int largeVertexSize;
	unsigned int distBound;
	int pairSupport;
	int resampling;
	bool inward;
	const char* inputName;
	const char* outputName;
};

/**
 * 	Get output file prefix with out extension name.
 *
 * @author	Tao Yu
 * @date	2011/3/29
 *
 * @param [in]	filename	Filename of the output file.
 *
 * @return	The full path with out extension name.
 */
string outFilePrefix(const string& filename)
{
	string prefix(filename);
	for (int i = 0; i < prefix.size(); ++i) {
		if (prefix[i] == '\\') {
			prefix[i] = '/';
		}
	}
	string::size_type lastdot = filename.rfind('.');
	string::size_type lastslash = filename.rfind('/');
	if (lastslash == string::npos || lastslash > lastdot) {
		return prefix;
	}
	return prefix.substr(0, lastdot + 1);
}

/**
 * 	Saves debuijn graph to disk.
 *
 * @author	Tao Yu
 * @date	2011/4/2
 *
 * @param [in,out]	dbgraph	If non-null, the debuijn graph.
 * @param	dataFile	   	The data file.
 * @param	signatureFile  	The signature file.
 *
 * @return	true if it succeeds, false if it fails.
 */
bool fsave( DeBruijnGraph* dbgraph, ReadInputFacts* facts, string dataFile, string signatureFile ) 
{
	// open files
	FILE *dataFp = fopen(dataFile.c_str(), "wb");
	if (!dataFp) {
		return false;
	}
	FILE*signatureFp = fopen(signatureFile.c_str(), "wb");
	if (!signatureFp) {
		return false;
	}
	if (!dbgraph->write_data(dataFp) || !fwrite(facts, sizeof(ReadInputFacts), 1, dataFp)) {
		fclose(dataFp);
		fclose(signatureFp);
		return false;
	}
	// checksum
	GINT64 fsize = ftello(dataFp);
	if (!fwrite(&fsize, sizeof(fsize), 1, signatureFp)) {
		fclose(dataFp);
		fclose(signatureFp);
		return false;
	}
	fclose(signatureFp);
	fclose(dataFp);
	return true;
}

/**
 * 	Validates if the debuijn graph data on disk is valid.
 * 	The validation is done by check if the file exists and size correct.
 *
 * @author	Tao Yu
 * @date	2011/4/2
 *
 * @param	dataFile	 	The data file.
 * @param	signatureFile	The signature file.
 *
 * @return	null if it fails, else.
 */
FILE* fvalidate( string dataFile, string signatureFile ) 
{
	// open files
	FILE *dataFp = fopen(dataFile.c_str(), "rb");
	if (!dataFp) {
		return NULL;
	}
	FILE*signatureFp = fopen(signatureFile.c_str(), "rb");
	if (!signatureFp) {
		return NULL;
	}
	if (fseeko(dataFp, 0ll, SEEK_END)) {
		fclose(dataFp);
		fclose(signatureFp);
		return NULL;
	}
	// checksum
	GINT64 fsize = ftello(dataFp);
	GINT64 sig;
	if (!fread(&sig, sizeof(sig), 1, signatureFp) || fsize != sig) {
		fclose(dataFp);
		fclose(signatureFp);
		return NULL;
	}
	fclose(signatureFp);
	fseeko(dataFp, 0ll, SEEK_SET);
	return dataFp;
}

/**
 * 	run. Excute from reading file to remove low coverage.
 *
 * @author	Tao Yu
 * @date	2011/3/23
 *
 * @param	kmerLength			 	Length of the kmer.
 * @param kmerThreshold		Minimal coverage of k-mer needed to be kept
 * @param	lengthThreshold		 	The length threshold.
 * @param	multiplicityThreshold	The multiplicity threshold.
 * @param	lowCoverage			 	The low coverage.
 * @param	insertLength		 	Length of the insert.
 * @param	insertLenSD			 	The insert length std var.
 * @param	disBound			 	The distance bound in remove bubbles.
 * @param pairSupport		 Minimal pair support required to establish trusted pair nodes.
 * @param resampling			The resampling rate of read file
 * @param	largeVertexSize		 	Size of the large vertex to be output.
 * @param	inputFasta			 	The input fasta.
 * @param	outputFasta			 	The output fasta.
 */
void run(ConfigArgs& confArgs)
{
	string prefix = outFilePrefix(string(confArgs.outputName));
	//string kmerCounterDataFile = "";
	string debruijnGraphDataFile = prefix + "dbgraph";
	string debruijnGraphValidationFile = prefix + "dbgraph.sig";
	//string sequenceGraphDataFile = "";
	string sequenceGraphDotFile = prefix + "seqg.dot";
	string sequenceGraphStatisFile = prefix + ".contig.statis";
	KmerVertex::setLength(confArgs.kmerLength); //< crucial step, do not forget
	
	DeBruijnGraph* dbgraph = NULL;
	ReadInputFacts* readFacts = NULL;
	if (FILE *dbgraphFp = fvalidate(debruijnGraphDataFile, debruijnGraphValidationFile))
	{
		cout << "Reading DeBuijn Graph data..." << endl;
		LOG(INFO) << "Reading DeBuijn Graph data...";
		dbgraph = new DeBruijnGraph();
		if (!dbgraph->read_data(dbgraphFp)) {
			LOG(FATAL) << "Reading DeBuijn Graph data failed.";
		}
		readFacts = new ReadInputFacts;
		if (!fread(readFacts, sizeof(ReadInputFacts), 1, dbgraphFp))
		{
			LOG(FATAL) << "Reading Read Facts data failed.";
		}
		cout << "Reading DeBuijn Graph data succeeded." <<endl;
		LOG(INFO) << "Reading DeBuijn Graph data succeeded.";
		fclose(dbgraphFp);
	}
	else
	{
		OneFilePairedFastaReader reader(confArgs.inputName);
		DeBruijnGraphBuilder builder(reader, confArgs.inward, confArgs.kmerThreshold);

		readFacts = builder.getReadFacts();
		GCOUNT_T numRead = readFacts->numOfRead();

		dbgraph = builder.getDeBruijnGraph();
		GCOUNT_T numKmer = dbgraph->numOfVertex();
		cout <<"Total number of k-mers: " <<numKmer <<endl;
		readFacts->setInsertionLength(confArgs.insertLength);
		readFacts->setInsertionStdDev(confArgs.insertLenSD);

		cout << "Saving DeBuijn Graph data..." << endl;
		LOG(INFO) << "Saving DeBuijn Graph data...";
		if (!fsave(dbgraph, readFacts, debruijnGraphDataFile, debruijnGraphValidationFile)) {
			cerr << "Saving DeBuijn Graph data failed." << endl;
			LOG(ERROR) << "Saving DeBuijn Graph data failed.";
		}
		cout << "Saving DeBuijn Graph data succeeded." <<endl;
		LOG(INFO) << "Saving DeBuijn Graph data succeeded.";
	}

	SequenceGraph seqgraph(dbgraph);
	delete dbgraph;
	dbgraph = NULL;
	GraphExaminator::exam(seqgraph);
	cout << "#Vertex: " << seqgraph.numOfVertex() <<endl;

	GraphErrorCorrector correctior(&seqgraph, confArgs.lengthThreshold, confArgs.multiplicityThreshold, confArgs.lowCoverage,confArgs.distBound, ceil(confArgs.distBound * 0.2), 5, 0.2);
	correctior.correctGraph();
	GraphExaminator::exam(seqgraph);

	// Output sequence graph
	cout << "Writing sequence graph to file." <<endl;
	LOG(INFO) << "Writing sequence graph to file.";
	SequenceGraphViz gvz(&seqgraph, prefix);
	gvz.writeData();
	cout << "Number of components: " << gvz.numOfComponents() << endl;
	cout << "Largest component: " << gvz.largestComponentSize() << " (" << setiosflags(ios_base::showpoint) << (float)gvz.largestComponentSize() * 100 / gvz.totalComponentSize() << "%)" <<endl ;
	cout << "Writing file succeeded." <<endl;
	LOG(INFO) << "Writing file succeeded.";

	//cout<< "begin do possible merge"<<endl;
	//seqgraph.doPossibleMerge();
	//cout<< "end do possible merge"<<endl;
	//cout << "#Vertex: " <<seqgraph.numOfVertex() <<endl;
	//int largeVertex = 0;

	//seqgraph.resetCursor();
	//fstream fout(outputFasta, ios::out);
	//if (!fout.is_open()) {
	//	cerr << "Error while writing result. (" << outputFasta << ")" <<endl;
	//}
	//SequenceVertex* v;
	//while (seqgraph.hasNext()) {
	//	v = seqgraph.next();
	//	if (v->m_numOfKmer > largeVertexSize) {
	//		fout << "> Contig " << ++largeVertex << endl;
	//		fout << v->decompressVertex(seqgraph, buffer) << endl;
	//		cout<<"length: "<<strlen(buffer)<<endl;
	//	}
	//}
	//fout.close();
	//cout << "Large vertex: " << largeVertex << endl;

	Disambiguation disambiguation(&seqgraph, *readFacts, confArgs.pairSupport, confArgs.resampling, confArgs.outputName, sequenceGraphStatisFile.c_str());
	disambiguation.allContigs();
	cout <<"Greatest contig length: " <<disambiguation.getLongestLength() <<endl;
	cout <<"Total number of contigs: " <<disambiguation.getNumOfContigs() <<endl;
	delete readFacts;
	readFacts = NULL;
}

/**
 * 	Show parameter list of the program.
 *
 * @author	Tao Yu
 * @date	2011/3/23
 */
void usage()
{
	cout <<"Sequence Assembly Program Usage:" <<endl;
	cout <<"-k kmerLength \tLength of kmer. Default: 21." <<endl;
	cout <<"-t kmerThreshold \tMinimal kmer coverage. Default: 2." <<endl;
	cout <<"-r lengthThreshold \tMinimal free ending length threshold. Default: 42." <<endl;
	cout <<"-m multiplicityThreshold \tMinimal multiplicity trusted. Default: 2." <<endl;
	cout <<"-c lowCoverage \tMinimal contig coverage. Default: 3." <<endl;
	cout <<"-l insertLength \tLibrary size. Default: 100." <<endl;
	cout <<"-v insertLenSD \tLibrary size standard variance. Default: 5." <<endl;
	cout <<"-b distBound \tHow far will the algorithm search to remove bubbles. Default: 200." <<endl;
	cout <<"-p minimal pair support required to establish trusted pair nodes. Default: 3" << endl;
	cout <<"-x the resampling rate of read file. Default: 50" <<endl;
	cout <<"-d direction of paired reads. Default: 1 (inward)" <<endl;
	cout <<"-s largeVertexSize \tSmallest contig to keep. Default: 200." <<endl;
	cout <<"-i inputName \tInput FASTA file. Default: ../examples/test_reads.fa." <<endl;
	cout <<"-o outputName \tOutput FASTA file. Default: ../examples/contig.fa." <<endl;
	cout <<"-h \t Show this help information." <<endl;
}

int main(int argc, char* argv[])
{
	// Initialize Google's logging library.
	google::InitGoogleLogging(argv[0]);

#ifdef __linux__
	google::InstallFailureSignalHandler();
#endif
	
	// variable defaults
	ConfigArgs confArgs;
	confArgs.kmerLength = 21;
	confArgs.kmerThreshold = 2;
	confArgs.multiplicityThreshold = 2;
	confArgs.lengthThreshold = 42;
	confArgs.lowCoverage = 3;
	confArgs.insertLength = 100;
	confArgs.insertLenSD = 5;
	confArgs.largeVertexSize = 200;
	confArgs.distBound = 200;
	confArgs.pairSupport = 3;
	confArgs.resampling = 50;
	confArgs.inward = true;
	confArgs.inputName = "../examples/test_reads.fa";
	confArgs.outputName = "../examples/test_reads.contig.fa";
	
	int ch;
	while((ch = getopt(argc, argv, "k:t:r:m:c:l:v:b:p:x:d:s:i:o:h"))!= -1)
	{
		switch(ch)
		{
		case 'k':
			confArgs.kmerLength = atoi(optarg);
			break;
		case 't':
			confArgs.kmerThreshold = atoi(optarg);
			break;
		case 'r':
			confArgs.lengthThreshold = atoi(optarg);
			break;
		case 'm':
			confArgs.multiplicityThreshold = atof(optarg);
			break;
		case 'c':
			confArgs.lowCoverage = atoi(optarg);
			break;
		case 'l':
			confArgs.insertLength =atoi(optarg);
			break;
		case 'v':
			confArgs.insertLenSD = atoi(optarg);
			break;
		case 'b':
			confArgs.distBound = atoi(optarg);
			break;
		case 'p':
			confArgs.pairSupport = atoi(optarg);
			break;
		case 'x':
			confArgs.resampling = atoi(optarg);
			break;
		case 'd':
			confArgs.inward = atoi(optarg);
			break;
		case 's':
			confArgs.largeVertexSize = atoi(optarg);
			break;
		case 'i':
			confArgs.inputName = optarg;
			break;
		case 'o':
			confArgs.outputName = optarg;
			break;
		case 'h':
			usage();
			return 0;
		case '?':
			// getopt already put error message
		default:
			usage();
			abort(); //return val = 134?
		}
	}

	cout <<"Current working directory: " <<endl;
	system("cd && pwd 2>>nul");
	cout << endl;

	cout <<confArgs <<endl;
	run(confArgs);

	cout <<"Assembly finished." <<endl;
	return 0;
}
